package cn.ai.shoes.client.service.impl;

import cn.ai.shoes.client.api.ComfyuiApi;
import cn.ai.shoes.client.entity.ComfyuiQueueTask;
import cn.ai.shoes.client.entity.CurrentServerEntity;
import cn.ai.shoes.client.service.ComfyuiService;
import cn.ai.shoes.client.utils.ComfyuiServerUtil;
import cn.ai.shoes.common.exception.ServiceException;
import cn.ai.shoes.common.exception.constant.GlobalErrorCodeConstants;
import cn.ai.shoes.core.pojo.entity.MyUser;
import cn.ai.shoes.core.service.UserService;
import cn.ai.shoes.core.utils.SecurityUtils;
import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.RandomUtil;

import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.toolkit.ObjectUtils;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import okhttp3.ResponseBody;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.stereotype.Component;
import retrofit2.Call;
import retrofit2.Response;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.io.FileOutputStream;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.*;

@Slf4j
@Component
public class ComfyuiServiceImpl implements ComfyuiService {
    @Resource
    private OkHttpClient okHttpClient;
    @Resource
    private ComfyuiServerUtil comfyuiServerUtil;
    @Value("${comfyui.server.url}")
    private String viewPaths;

    // 每次动态获取最新服务器API实例
    private ComfyuiApi getComfyuiApi(String serverUrl) {
        Retrofit retrofit = new Retrofit.Builder()
                .baseUrl(serverUrl)
                .client(okHttpClient)
                .addConverterFactory(JacksonConverterFactory.create())
                .build();
        return retrofit.create(ComfyuiApi.class);
    }

    /**
     * 获取图片路径
     *
     * @param filename
     * @param type
     * @param subfolder
     * @return
     */
    @NotNull
    @Override
    public String getViewUrl(String serverId,String filename, String type, String subfolder) {
        String viewUrl = "http://"+viewPaths+"/file/view?serverId="+serverId+"&filename=" + filename + "&type=" + type + "&subfolder=" + subfolder;
        return viewUrl;
    }

    /**
     * 回显视图
     *
     * @return
     */
    @Override
    public ResponseEntity<byte[]> getView(String serverId,String filename, String type, String subfolder) {
        //获取当前服务器地址列表
        String urlByKey = comfyuiServerUtil.getUrlByKey(serverId);
        //如果urlByKey为空，则返回错误信息
        if (ObjectUtils.isEmpty(urlByKey)) {
            throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_SERVER_NULL_EXIST);
        }
        Call<ResponseBody> viewResult = getComfyuiApi(urlByKey).getView(filename, type, subfolder);
        try {
            Response<ResponseBody> executeViewResult = viewResult.execute();
            if (!executeViewResult.isSuccessful()) {
                // comfyui返回图片失败
                throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_FILE_NULL);
            }
            ResponseBody viewResultBody = executeViewResult.body();
            if (viewResultBody == null) {
                throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_API_RETURN_NULL);
            }
            byte[] imageBytes = viewResultBody.bytes();
            return ResponseEntity.ok()
                    .contentType(MediaType.IMAGE_PNG)
                    .body(imageBytes);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取图片路径
     * @param historyResultBody
     * @param promptId
     * @return
     */
    @NotNull
    public List<String> getViewUrlList(HashMap historyResultBody, String promptId,String currentServerUrl) {
        List<String> viewUrlList = new ArrayList();
        //获取图片信息
        Map<String,Object> promptBody = Convert.toMap(String.class, Object.class,historyResultBody.get(promptId));
        //获取outputs信息
        Map<String,?> outputs= BeanUtil.toBean(promptBody.get("outputs"),Map.class);
        //获取images信息
        outputs.forEach((key,value)->{
            HashMap imagesBody= (HashMap) value;
            //判断imagesBody.get("images")是否为空
            if (imagesBody.get("images") ==null){
                return;
            }
            ArrayList images = (ArrayList) imagesBody.get("images");
                for (Object image : images) {
                HashMap imageBody= (HashMap) image;
                String filename = imageBody.get("filename").toString();
                String subfolder = imageBody.get("subfolder").toString();
                String type = imageBody.get("type").toString();
                //获取当前服务器地址列表
                String serverId = comfyuiServerUtil.getKeyByUrl(currentServerUrl);
                //拼接图片路径为get请求后缀格式
                String viewUrl = getViewUrl(serverId,filename, type, subfolder);
                //将comfyui图片保存到本地
                // comfyuiLocalSaveImage(viewUrl,filename,type,subfolder);
                //将图片路径添加到url数组中
                    viewUrlList.add(viewUrl);
            }
        });
        return viewUrlList;
    }

    /**
     * 保存图片到本地
     * @param viewUrl
     * @param filename
     * @param type
     * @param subfolder
     */
    public void comfyuiLocalSaveImage(String viewUrl, String filename, String type, String subfolder) {
        //调用comfyui api接口获取预览的图片信息
        Call<ResponseBody> viewResult = getComfyuiApi(viewUrl).getView(filename, type, subfolder);
        try {
            Response<ResponseBody> executeViewResult = viewResult.execute();
            if (!executeViewResult.isSuccessful()){
                //comfyui保存图片失败
                throw new ServiceException(GlobalErrorCodeConstants.Mask_ERROR.getResult());
            }
            ResponseBody viewResultBody = executeViewResult.body();
            if (viewResultBody!=null){
                //当前日期+小时+分钟+秒+随机三个字母+.png
                String savaFileName= new SimpleDateFormat("yyyyMMddHHmmss").format(new Date())+ RandomUtil.randomString(3)+".png";
                //保存图片到本地
                viewResultBody.byteStream().transferTo(new FileOutputStream("D:\\comfyui-image\\output\\"+savaFileName));
            }
        } catch (IOException e) {
            //comfyui保存图片失败
            throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_FAILED_SAVE_IMAGE.getResult());
        }
    }

    /**
     * 等待生成图片任务完成
     * @param
     * @return
     */
    @NotNull
    public HashMap waitingTask(String promptId,String currentServerUrl) throws IOException, InterruptedException {
        HashMap historyResultBody = null;
        //等待5分钟
        for (int i = 0; i < 300; i++) {
            //使用slf4j打印
            log.info("调用comfyui返回图片中===============第"+i+"轮");
            //等待任务完成
            Call<HashMap> historyResult = getComfyuiApi(currentServerUrl).getHistoryTask(promptId);
            Response<HashMap> executeHistory;
            executeHistory = historyResult.execute();
            if (!executeHistory.isSuccessful()){
                throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_API_ERROR.getResult());
            }
            historyResultBody = executeHistory.body();
            //判断任务是否完成
            if (historyResultBody.get(promptId)!=null){
                //任务完成
                break;
            }
            if (i==299){
                throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_TIMEOUT_ERROR.getResult());
            }
            //休息1秒
            Thread.sleep(1000);

        }
        return historyResultBody;
    }
    /**
     * 发起任务
     * @param prompt
     * @return
     */
    public String generateTask(JSONObject prompt,String currentServerUrl) throws IOException {
        //从springSecurityContext中获取用户信息
        MyUser myUser = SecurityUtils.getCurrentUser();
        if (ObjectUtils.isEmpty(myUser)){
            throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_USER_NOT_LOGIN.getResult());
        }
        ComfyuiQueueTask comfyuiQueTask= ComfyuiQueueTask.builder()
                .clientId(myUser.getUid())
                .prompt(prompt)
                .build();
        Response<HashMap> promptIdResult = generateImages(comfyuiQueTask,currentServerUrl);
        if (!promptIdResult.isSuccessful()) {
            throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_API_ERROR.getResult());
        }
        HashMap promptIdResultBody = promptIdResult.body();
        String promptId = promptIdResultBody.get("prompt_id").toString();
        return promptId;
    }

    /**
     * 调用comfyui生成图片
     * @param comfyuiQueTask
     * @return
     */
    @NotNull
    public Response<HashMap> generateImages(ComfyuiQueueTask comfyuiQueTask,String currentServerUrl) throws IOException {
        //调用comfyui api接口
        Call<HashMap> promptIdCall = getComfyuiApi(currentServerUrl).addQueueTask(comfyuiQueTask);
        Response<HashMap> promptIdResult = promptIdCall.execute();
        //判断请求是否成功
        if (!promptIdResult.isSuccessful()){
            // 返回异常信息
            throw new ServiceException(GlobalErrorCodeConstants.COMFYUI_API_RETURN_NULL);
        }
        return promptIdResult;
    }
}
